from typing import Any, Awaitable, Callable, Dict

from aiogram import BaseMiddleware
from aiogram.types import TelegramObject
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from tinkoff_bot.infrastructure.database.services import UserService
from tinkoff_bot.infrastructure.database.uow import build_uow


class DBSessionMiddleware(BaseMiddleware):
    def __init__(self, session_maker: async_sessionmaker[AsyncSession]) -> None:
        self.session_maker = session_maker

    async def __call__(
        self,
        handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
        event: TelegramObject,
        data: Dict[str, Any],
    ) -> Any:
        async with self.session_maker() as session:
            uow = build_uow(session)
            data["uow"] = uow
            data["user_service"] = UserService(uow)

            result = await handler(event, data)
            del data["uow"]
        return result
